from cio.jsonl import JsonLReader, JsonLWriter
from config.Config import BASE_DIR
from utils.splitter import code_split, nl_split

train_data = JsonLReader(BASE_DIR / 'resource' / 'data' / 'consum' / 'context.train.jsonl').yield_read()
valid_data = JsonLReader(BASE_DIR / 'resource' / 'data' / 'consum' / 'context.test.jsonl').yield_read()

#### 使用环境
context_src_train_file = open(BASE_DIR / 'resource' / 'data' / 'consum' / 'context-src.txt', "w+", encoding="utf-8")
context_src_valid_file = open(BASE_DIR / 'resource' / 'data' / 'consum' / 'context-src-valid.txt', "w+",
                              encoding="utf-8")

src_train_file = open(BASE_DIR / 'resource' / 'data' / 'consum' / 'src-train.txt', "w+", encoding="utf-8")
src_valid_file = open(BASE_DIR / 'resource' / 'data' / 'consum' / 'src-valid.txt', "w+",
                      encoding="utf-8")

tgt_train_file = open(BASE_DIR / 'resource' / 'data' / 'consum' / 'tgt-train.txt', "w+", encoding="utf-8")
tgt_valid_file = open(BASE_DIR / 'resource' / 'data' / 'consum' / 'tgt-valid.txt', "w+", encoding="utf-8")

for item in train_data:
    method = " ".join(code_split(item['method']))
    summary = " ".join(nl_split(item['summary']))
    context = " ".join(['<pn>'] + code_split(item['package_name']) + ['</pn>'])
    context += " ".join(['<cs>'] + code_split(item['class_signature']) + ['</cs>'])
    context_src = method + " " + context
    src = method
    tgt = summary
    context_src_train_file.write(context_src + "\n")
    src_train_file.write(src + "\n")
    tgt_train_file.write(tgt + "\n")


for item in valid_data:
    method = " ".join(code_split(item['method']))
    summary = " ".join(nl_split(item['summary']))
    context = " ".join(['<pn>'] + code_split(item['package_name']) + ['</pn>'])
    context += " ".join(['<cs>'] + code_split(item['class_signature']) + ['</cs>'])
    context_src = method + " " + context
    src = method
    tgt = summary
    context_src_valid_file.write(context_src + "\n")
    src_valid_file.write(src + "\n")
    tgt_valid_file.write(tgt + "\n")